# -*- coding: utf-8 -*-
'''
Created on Sep 2, 2019

@author: yl
'''
from PyUeye.Full_Frame import file_name, img_name, doc_name, txt_name
import numpy as np
import os
import matplotlib.pyplot as plt
import cv2
from scipy import signal
from FFT_Interpolation import FFT_interpolation_2, line_cal
from scipy.optimize import curve_fit
from scipy import signal
import datetime
now = datetime.datetime.now()

def Export_Data(file_name, header, out_str):
    print('Writing Data')
    with open(file_name,'w') as fid: ######################################################################################
        fid.writelines(header)
        fid.writelines(out_str)
    print('Finish Writing')
    return

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q
def fit_func_line(x, p, q):
    return p*x+q

#############
# fs_cam = 575.8  ###USB2.0
# fs_cam = 1370.72 ###USB3.0
# fs_cam = 1860 ###USB3.0 numpy readout
fs_cam = 50 ### Hardware Trigger
####################

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_f_fit_set = []
ver_f_fit_set = []
hor_phase_centers = []
ver_phase_centers = []

f = open(img_name, 'rb')
f_size = os.fstat(f.fileno()).st_size
frames = []
while f.tell() < f_size:
#     print(f.tell())
    frame = np.load(f, allow_pickle=True)
    frames.append(frame)

# x_axis = np.linspace(0, len(time_sequence)-1, len(time_sequence))

frames = np.array(frames[0])
print(frames.shape)

if len(frames.shape) == 3:
    print('file_size: %.1fMB'%(f_size/1e6))
    print('%i frames'%len(frames))
    print('%i line/frame'%len(frames[0]))
    print('%i pixel/line'%len(frames[0][0]))

print('frame rate: %.2f'%fs_cam)
print('record time: %f s'%(len(frames)/fs_cam))


img_set = frames
frame_num = len(frames)

hor_threshold, ver_threshold = 100, 100
hor_centerline, ver_centerline = img_set[int(frame_num/2)][int(len(img_set[0])/2), :], img_set[int(frame_num/2)][:, int(len(img_set[0][1])/2)]
hor_start, ver_start = np.where(hor_centerline > hor_threshold)[0][0], np.where(ver_centerline > ver_threshold)[0][0]
hor_end, ver_end = np.where(hor_centerline > hor_threshold)[0][-1], np.where(ver_centerline > ver_threshold)[0][-1]
hor_index = int((ver_start+ver_end)/2)
ver_index = int((hor_start+hor_end)/2)

num = 0
line_num = 100
for img in img_set[::]:
    num += 1
    hor_phases, ver_phases = [], []
    hor_freqs, ver_freqs = [], []
    for i in range(line_num):
        hor_center = np.array(img[int(hor_index-line_num/2+i)])
        ver_center = img[:,int(ver_index-line_num/2+i)]
        
        hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
        ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal(ver_center)

        hor_freqs.append(hor_f_fit)
        hor_phases.append(hor_phase_estim)
        ver_freqs.append(ver_f_fit)
        ver_phases.append(ver_phase_estim)
    
    hor_phases = np.array(hor_phases)
    hor_freqs = np.array(hor_freqs)
    ver_phases = np.array(ver_phases)
    ver_freqs = np.array(ver_freqs)
    
    hor_f_fit_set.append(np.sum(hor_freqs)/line_num)
    ver_f_fit_set.append(np.sum(ver_freqs)/line_num)
    
    fitted_phase = np.sum(np.unwrap(hor_phases))/line_num
    if fitted_phase > np.pi:
        fitted_phase = fitted_phase - 2*np.pi
    if fitted_phase < -np.pi:
        fitted_phase = fitted_phase + 2*np.pi
    hor_phase_centers.append(fitted_phase)
    
    fitted_phase = np.sum(np.unwrap(ver_phases))/line_num
    if fitted_phase > np.pi:
        fitted_phase = fitted_phase - 2*np.pi
    if fitted_phase < -np.pi:
        fitted_phase = fitted_phase + 2*np.pi
    ver_phase_centers.append(fitted_phase)
    print(num, hor_m_k_num, ver_m_k_num)

hor_f_fit_set = np.array(hor_f_fit_set)
ver_f_fit_set = np.array(ver_f_fit_set)
hor_phase_centers = np.array(hor_phase_centers)
ver_phase_centers = np.array(ver_phase_centers)

hor_angle = (V_x-Lamda*hor_f_fit_set/2)*1e6 ### urad
hor_length = np.unwrap(hor_phase_centers)/4/np.pi*Lamda*1e9 ### nm
ver_angle = (V_y-Lamda*ver_f_fit_set/2)*1e6 ### urad
ver_length = np.unwrap(ver_phase_centers)/4/np.pi*Lamda*1e9 ### nm

''' 
    输出数据
'''

header = ['%s\n' %doc_name,
      'Local current time : %s\n' %now.strftime("%Y-%m-%d %H:%M:%S"),
      'Fs = %e (Hz)\n' %fs_cam,##########################################################################################################
      'Record time: %e (s)\n' %(len(frames)/fs_cam),############################################################################################
      'Frame Number = %i\n' %len(frames),############################################################################################
      'Lines/Frame = %i\n' %len(frames[0]),############################################################################################
      'Pixels/Line = %i\n' %len(frames[0][0]),############################################################################################
      'Channel_1: Horizontal freq.(16-line averaged)\n',############################################################################################
      'Channel_2: Horizontal phase(16-line averaged)\n',############################################################################################
      'Channel_3: Vertical freq.(16-line averaged)\n',############################################################################################
      'Channel_4: Vertical phase(16-line averaged)\n',############################################################################################
#       'Channel_3: Demodulated S_curve (from CH_No.2)\n',############################################################################################
#       'Channel_4: Demodulated S_curve (from CH_No.1)\n',############################################################################################
#       'Data_Ch5: Difference between Channel_3 & Channel_4 (Ch3-Ch4)\n'################################################################################
      '-------------------------------------------------\n',
      ]
out_str = ['%f, %f, %f, %f\n' %(hor_f_fit_set[i], hor_phase_centers[i], ver_f_fit_set[i], ver_phase_centers[i]) for i in range(len(hor_f_fit_set))]    
 
''' 
    保存文件
'''
Export_Data(txt_name, header, out_str)
print('TXT file saved')



# plt.figure(1)
# im = plt.imshow(img_set[0], cmap='gray')
# plt.colorbar(im, fraction=0.046, pad=0.04)
# plt.show()
# plt.figure(2)
# plt.plot(ver_f_fit_set)
# plt.show()

if 1:### Plotting Noise
    ######Noise_PSD#### linearity range
    start_cutting, stop_cutting = 0, len(hor_angle)
#     start_cutting, stop_cutting = 20, 3000
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting) 
    ############### Plotting
     
    plt.figure('Tilt')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (hor_angle-np.average(hor_angle))[range[0]:range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (ver_angle-np.average(ver_angle))[range[0]:range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle, fs_cam, nperseg=len(hor_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise")
    plt.Noise_PSDl("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_angle, fs_cam, nperseg=len(ver_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Tilting Noise")
    plt.Noise_PSDl("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    
    plt.figure('Length')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (hor_length-np.average(hor_length))[range[0]:range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (ver_length-np.average(ver_length))[range[0]:range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_length, fs_cam, nperseg=len(hor_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise")
    plt.Noise_PSDl("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_length, fs_cam, nperseg=len(ver_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Length Noise")
    plt.Noise_PSDl("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    
    plt.show()

    
if 0:### Plotting Nonlinearity
    ############### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
    start_cutting, stop_cutting = start_cutting/fs_cam, stop_cutting/fs_cam
    zoom_time = start_cutting, stop_cutting
    
    start_cutting, stop_cutting = 1.6, 6.2 ### in second   
#     zoom_time = [1.2, 1.4]
    
    start_cutting, stop_cutting = int(start_cutting*fs_cam), int(stop_cutting*fs_cam)
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
    selected_time = (np.arange(stop_cutting-start_cutting)+start_cutting)/fs_cam    
    zoom_range = [zoom_time[0]*fs_cam, zoom_time[1]*fs_cam]
     
    plt.figure('Tilt')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_angle = hor_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    hor_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_angle = ver_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    ver_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    plt.plot(x, y, color='red', marker=' ')
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

#     ax5 = plt.subplot(4,2,5)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = hor_angle_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Horizontal Tilting")
#     plt.ylabel("Horizontal Angle (urad)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')
#        
#     ax6 = plt.subplot(4,2,6)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = ver_angle_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Vertical Tilting")
#     plt.ylabel("Vertical Angle (urad)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')

    ax7 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle_nonlinearity, fs_cam, nperseg=len(hor_angle_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise")
    plt.Noise_PSDl("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax8 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_angle_nonlinearity, fs_cam, nperseg=len(ver_angle_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Tilting Noise")
    plt.Noise_PSDl("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    
    
    #### Length
    plt.figure('Length')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_length = hor_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    hor_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_length = ver_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    ver_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    plt.plot(x, y, color='red', marker=' ')
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
    
#     ax5 = plt.subplot(4,2,5)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = hor_length_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Horizontal Length")
#     plt.ylabel("Horizontal Length (nm)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')
#        
#     ax6 = plt.subplot(4,2,6)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = ver_length_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Vertical Length")
#     plt.ylabel("Vertical Length (nm)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')

    ax7 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_length_nonlinearity, fs_cam, nperseg=len(hor_length_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise")
    plt.Noise_PSDl("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-3, 2e3) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax8 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_length_nonlinearity, fs_cam, nperseg=len(ver_length_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Length Noise")
    plt.Noise_PSDl("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-3, 2e3) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    plt.show()